#!/bin/bash

# BERT意图识别服务启动脚本

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    if [ "$DEBUG" = "true" ]; then
        echo -e "${BLUE}[DEBUG]${NC} $1"
    fi
}

# 显示启动信息
show_banner() {
    echo "========================================"
    echo "🤖 BERT意图识别服务"
    echo "========================================"
    echo "版本: 1.0.0"
    echo "Python: $(python --version)"
    echo "工作目录: $(pwd)"
    echo "时间: $(date)"
    echo "========================================"
}

# 检查环境
check_environment() {
    log_info "检查运行环境..."
    
    # 检查Python
    if ! command -v python &> /dev/null; then
        log_error "Python未安装"
        exit 1
    fi
    
    # 检查必要的目录
    for dir in logs models data .cache; do
        if [ ! -d "$dir" ]; then
            log_warn "创建目录: $dir"
            mkdir -p "$dir"
        fi
    done
    
    # 检查权限
    if [ ! -w "logs" ]; then
        log_warn "logs目录无写权限"
    fi
    
    log_info "环境检查完成"
}

# 检查模型文件
check_model() {
    local model_path="${MODEL_PATH:-./final-bert-intent-model}"
    
    log_info "检查模型文件: $model_path"
    
    if [ ! -d "$model_path" ]; then
        log_warn "模型目录不存在: $model_path"
        log_warn "请确保模型已正确挂载或复制到容器中"
        return 1
    fi
    
    # 检查关键文件
    local required_files=("config.json" "pytorch_model.bin" "tokenizer_config.json")
    for file in "${required_files[@]}"; do
        if [ ! -f "$model_path/$file" ]; then
            log_warn "缺少模型文件: $model_path/$file"
        fi
    done
    
    log_info "模型检查完成"
    return 0
}

# 初始化配置
init_config() {
    log_info "初始化配置..."
    
    # 设置默认环境变量
    export PYTHONPATH="${PYTHONPATH:-/app}"
    export HOST="${HOST:-0.0.0.0}"
    export PORT="${PORT:-5000}"
    export MODEL_PATH="${MODEL_PATH:-./final-bert-intent-model}"
    export LOG_LEVEL="${LOG_LEVEL:-INFO}"
    export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-/app/.cache}"
    export HF_HOME="${HF_HOME:-/app/.cache}"
    
    log_debug "PYTHONPATH: $PYTHONPATH"
    log_debug "HOST: $HOST"
    log_debug "PORT: $PORT"
    log_debug "MODEL_PATH: $MODEL_PATH"
    log_debug "LOG_LEVEL: $LOG_LEVEL"
    
    log_info "配置初始化完成"
}

# 等待依赖服务
wait_for_dependencies() {
    log_info "等待依赖服务..."
    
    # 如果有Redis配置，等待Redis
    if [ -n "$REDIS_HOST" ]; then
        log_info "等待Redis服务: $REDIS_HOST:${REDIS_PORT:-6379}"
        while ! nc -z "$REDIS_HOST" "${REDIS_PORT:-6379}"; do
            log_warn "Redis服务未就绪，等待..."
            sleep 2
        done
        log_info "Redis服务已就绪"
    fi
    
    log_info "依赖服务检查完成"
}

# 启动API服务
start_api() {
    log_info "启动BERT意图识别FastAPI服务..."
    log_info "服务地址: http://$HOST:$PORT"
    log_info "API文档: http://$HOST:$PORT/docs"
    log_info "ReDoc文档: http://$HOST:$PORT/redoc"
    
    # 检查模型（非严格检查）
    check_model || log_warn "模型检查失败，服务可能无法正常工作"
    
    # 设置uvicorn参数
    local workers="${WORKERS:-1}"
    local log_level="${LOG_LEVEL,,}"  # 转换为小写
    
    # 开发模式参数
    local reload_flag=""
    if [ "$DEBUG" = "true" ]; then
        reload_flag="--reload"
        log_info "开发模式：启用热重载"
    fi
    
    log_info "启动参数: workers=$workers, log_level=$log_level"
    
    # 启动FastAPI应用
    exec uvicorn src.api.app:app \
        --host "$HOST" \
        --port "$PORT" \
        --workers "$workers" \
        --log-level "$log_level" \
        --access-log \
        $reload_flag
}

# 运行训练
run_training() {
    log_info "启动模型训练..."
    
    local data_file="${DATA_FILE:-data/sample/intent_data.csv}"
    local config_file="${CONFIG_FILE:-configs/default_config.yaml}"
    local output_dir="${OUTPUT_DIR:-models/checkpoints}"
    
    log_info "数据文件: $data_file"
    log_info "配置文件: $config_file"
    log_info "输出目录: $output_dir"
    
    # 检查数据文件
    if [ ! -f "$data_file" ]; then
        log_error "数据文件不存在: $data_file"
        exit 1
    fi
    
    # 运行训练脚本
    exec python scripts/train.py \
        --data "$data_file" \
        --config "$config_file" \
        --output-dir "$output_dir" \
        --log-level "$LOG_LEVEL"
}

# 运行评估
run_evaluation() {
    log_info "启动模型评估..."
    
    local model_path="${MODEL_PATH:-./final-bert-intent-model}"
    local data_file="${DATA_FILE:-data/sample/intent_data.csv}"
    
    log_info "模型路径: $model_path"
    log_info "数据文件: $data_file"
    
    # 检查模型和数据
    if [ ! -d "$model_path" ]; then
        log_error "模型目录不存在: $model_path"
        exit 1
    fi
    
    if [ ! -f "$data_file" ]; then
        log_error "数据文件不存在: $data_file"
        exit 1
    fi
    
    # 运行评估脚本
    exec python scripts/evaluate.py \
        --model "$model_path" \
        --data "$data_file" \
        --log-level "$LOG_LEVEL"
}

# 运行预测
run_prediction() {
    log_info "启动批量预测..."
    
    local model_path="${MODEL_PATH:-./final-bert-intent-model}"
    local input_file="${INPUT_FILE:-data/input.txt}"
    local output_file="${OUTPUT_FILE:-data/output.json}"
    
    log_info "模型路径: $model_path"
    log_info "输入文件: $input_file"
    log_info "输出文件: $output_file"
    
    # 运行预测脚本
    exec python scripts/predict.py \
        --model "$model_path" \
        --input "$input_file" \
        --output "$output_file" \
        --log-level "$LOG_LEVEL"
}

# 显示帮助信息
show_help() {
    echo "用法: $0 [COMMAND]"
    echo ""
    echo "可用命令:"
    echo "  api        启动API服务 (默认)"
    echo "  train      运行模型训练"
    echo "  eval       运行模型评估"
    echo "  predict    运行批量预测"
    echo "  bash       启动bash shell"
    echo "  help       显示此帮助信息"
    echo ""
    echo "环境变量:"
    echo "  HOST              API服务主机 (默认: 0.0.0.0)"
    echo "  PORT              API服务端口 (默认: 5000)"
    echo "  MODEL_PATH        模型路径 (默认: ./final-bert-intent-model)"
    echo "  LOG_LEVEL         日志级别 (默认: INFO)"
    echo "  DEBUG             调试模式 (默认: false)"
    echo "  DATA_FILE         数据文件路径"
    echo "  CONFIG_FILE       配置文件路径"
    echo "  OUTPUT_DIR        输出目录"
    echo ""
}

# 主函数
main() {
    # 显示启动信息
    show_banner
    
    # 检查环境
    check_environment
    
    # 初始化配置
    init_config
    
    # 等待依赖服务
    wait_for_dependencies
    
    # 解析命令
    local command="${1:-api}"
    
    case "$command" in
        "api")
            start_api
            ;;
        "train")
            run_training
            ;;
        "eval"|"evaluate")
            run_evaluation
            ;;
        "predict")
            run_prediction
            ;;
        "bash"|"shell")
            log_info "启动bash shell..."
            exec /bin/bash
            ;;
        "help"|"--help"|"-h")
            show_help
            ;;
        *)
            log_error "未知命令: $command"
            show_help
            exit 1
            ;;
    esac
}

# 信号处理
trap_handler() {
    log_info "接收到停止信号，正在关闭服务..."
    exit 0
}

# 注册信号处理器
trap trap_handler SIGTERM SIGINT

# 运行主函数
main "$@" 